#include "header.h"
#include <linux/slab.h>
#define ORDER 5
#define SET_RELEASE_FN(fn) __releaseKeyFn = fn;
#define RELEASE_KEY(addr) __releaseKeyFn(addr);
#define BIGGER(n,m) n->tree->is_bigger(n,m)
#define SMALLER(n,m) n->tree->is_smaller(n,m) 
#define EQUAL(n,m) n->tree->is_equal(n,m) 
#define NOBIGGER(n,m) n->tree->is_smaller(n,m)||n->tree->is_equal(n,m)
#define ALLOW_NEW_KEY(n) (n->ksz<ORDER)
#define ALLOW_NEW_CHD(n) (n->csz<ORDER+1)
#define NEED_SPLIT(n) (n->ksz==ORDER)
#define HAS_SURPLUS_KEY(n) (n->ksz>(ORDER>>1))
#define NEED_KEY(n) (n->ksz<(ORDER>>1))
#define ALLOW_ROTATE_LEFT(n) (n && n->preSib && n->preSib->ksz<ORDER-1)
#define ALLOW_ROTATE_RIGHT(n) (n && n->nextSib && n->nextSib->ksz<ORDER-1)

/************************ BTree Logic *************************/
typedef struct BNode {
	ssize_t idx;
	ssize_t ksz;
	ssize_t csz;
	void** keys;
	struct BNode *parent, *preSib, *nextSib;
	struct BNode** children;	
	SelBalanceTree* tree;
} BNode, *BTree;

typedef struct QueryRet {
	BNode* node;
	void** key;
} QueryRet;

void __releaseBNode(BNode* node);
void __removeKeyFromNode(BNode*, void**, BNode**);
void __insertKeyToBNode(BNode*, void*, BNode*, BNode**);

BNode* __allocBNode(void* tree)
{
	BNode* node = (BNode*)kmalloc(sizeof(BNode), GFP_KERNEL);
	memset(node, 0, sizeof(BNode));
	node->tree = tree;
	node->keys = (void**)kmalloc(sizeof(void*)*ORDER, GFP_KERNEL);
	memset(node->keys, 0, sizeof(void*)*ORDER);
	node->children = (BNode**)kmalloc(sizeof(BNode*)*(ORDER+1), GFP_KERNEL);
	memset(node->children, 0, sizeof(BNode*)*(ORDER+1));
	return node;
}

void __rightMoveKeys(BNode* node) 
{
	void** lastPos = node->keys + node->ksz - 1;
	while (lastPos >= node->keys) {
		*(lastPos+1) = *lastPos;
		--lastPos;
	}
}

void __leftMoveKeys(BNode* node) 
{
	void** lastPos = node->keys;
	while (lastPos < node->keys+node->ksz-1) {
		*lastPos = *(lastPos+1);
		++lastPos;
	}
}

void __rightMoveChildren(BNode* node) 
{
	BNode** lastPos = node->children + node->csz - 1;
	while (lastPos >= node->children) {
		*(lastPos+1) = *lastPos;
		(*(lastPos+1))->idx += 1;
		--lastPos;
	}
}

void __leftMoveChildren(BNode* node) 
{
	BNode** lastPos = node->children;
	while (lastPos < node->children+node->csz-1) {
		*lastPos = *(lastPos+1);
		(*lastPos)->idx -= 1;
		++lastPos;
	}
	if (*node->children) node->children[0]->preSib = NULL;
}

void __rightRotateKey(BNode* node)
{
	void *lastKey, *parentKey;
	BNode *parent, *nextSib;
	parent = node->parent;
	nextSib = node->nextSib;
	
	lastKey = node->keys[node->ksz-1];	
	--node->ksz;
	
	parentKey = parent->keys[node->idx];
	parent->keys[node->idx] = lastKey;
	
	__rightMoveKeys(nextSib);
	nextSib->keys[0] = parentKey;
	++nextSib->ksz;
	
	if (node->csz) {
		BNode* lastChd = node->children[node->csz-1];
		--node->csz;
		if (node->csz) node->children[node->csz-1]->nextSib = NULL;
		__rightMoveChildren(nextSib);
		nextSib->children[0] = lastChd;
		lastChd->parent = nextSib;
		lastChd->preSib = NULL;
		lastChd->idx = 0;
		++nextSib->csz;
		if (nextSib->csz-1) {
			lastChd->nextSib = nextSib->children[1];
			nextSib->children[1]->preSib = lastChd;
		} else lastChd->nextSib = NULL;
	}
}

void __leftRotateKey(BNode* node)
{
	void *firstKey, *parentKey;
	BNode *parent, *preSib;
	parent = node->parent;
	preSib = node->preSib;
	
	firstKey = node->keys[0];	
	__leftMoveKeys(node);
	--node->ksz;
	
	parentKey = parent->keys[node->idx-1];
	parent->keys[node->idx-1] = firstKey;
	
	preSib->keys[preSib->ksz] = parentKey;
	++preSib->ksz;
	
	if (node->csz) {
		BNode* firstChd = node->children[0];
		__leftMoveChildren(node);
		--node->csz;
		preSib->children[preSib->csz] = firstChd;
		firstChd->parent = preSib;
		firstChd->idx = preSib->csz;
		firstChd->nextSib = NULL;
		if (firstChd->idx) {
			firstChd->preSib = preSib->children[firstChd->idx-1];
			preSib->children[firstChd->idx-1]->nextSib = firstChd;
		} else firstChd->preSib = NULL;
		++preSib->csz;
	}
}

BNode* __moveHalfKeyToNewRightSib(BNode* node)
{
	ssize_t idx, __idx;
	BNode* newNextSib = NULL;
	newNextSib = __allocBNode(node->tree);
	newNextSib->nextSib = node->nextSib;
	if (newNextSib->nextSib) newNextSib->nextSib->preSib = newNextSib;
	node->nextSib = newNextSib;
	newNextSib->preSib = node;
	
	idx = node->ksz>>1;
	memcpy(newNextSib->keys, node->keys+idx+1, (node->ksz-idx-1)*sizeof(void*));
	newNextSib->ksz += node->ksz - idx - 1;
	node->ksz -= node->ksz - idx - 1;
	
	if (node->csz) {
		memcpy(newNextSib->children, node->children+idx+1, (node->csz-idx-1)*sizeof(BTree));
		newNextSib->csz += node->csz-idx-1;
		node->csz -= node->csz-idx-1;
		__idx = 0;
		while (__idx < newNextSib->csz) {
			newNextSib->children[__idx]->idx = __idx;
			newNextSib->children[__idx]->parent = newNextSib;
			++__idx;
		}
		newNextSib->children[0]->preSib = NULL;
		node->children[node->csz-1]->nextSib = NULL;
	}
	
	return newNextSib;
}

BNode* __newRoot(void* key, BNode* leftChd, BNode* rightChd) 
{
	BNode* root = NULL;
	root = __allocBNode(leftChd?leftChd->tree:NULL);
	root->keys[0] = key;
	root->ksz = 1;
	
	if (leftChd && rightChd) {
		root->csz = 2;
		root->children[0] = leftChd;
		leftChd->idx = 0;
		root->children[1] = rightChd;
		rightChd->idx = 1;
		leftChd->parent = rightChd->parent = root;
	}
	
	return root;
}

void __spliteBNode(BNode* node, BNode** root)
{
	void* lastKey = NULL;
	BNode* parent = NULL;
	BNode* newNextSib = NULL;
	
	parent = node->parent;
	newNextSib = __moveHalfKeyToNewRightSib(node);
	lastKey = node->keys[node->ksz-1];
	node->keys[node->ksz-1] = NULL;
	--node->ksz;
	if (parent) __insertKeyToBNode(parent, lastKey, newNextSib, root);
	else *root = __newRoot(lastKey, node, newNextSib);
}

void** __findFirstBigger(BNode* node, void* key) 
{
	void **st, **ed, **mid;
	if (!node || !key) return NULL;
	st = node->keys;
	ed = st + node->ksz - 1;
	mid = NULL;
	while (st <= ed) {
		mid = st + ((ed-st)>>1);
		if (node->tree->is_equal(*mid, key) || 
		    node->tree->is_smaller(*mid, key)) st = mid + 1;
		else {
			if (ed == mid) break;
			else ed = mid;
		}
	}
	if (node->tree->is_bigger(*ed, key)) return ed;
	else return NULL;
}

void* GetFirstKeyFromBTree(BNode* root)
{
	if (!root) return NULL;
	return root->keys[0];
}

void QueryKeyInBTree(BNode* root, void* key, QueryRet* ret) 
{
	void** target = NULL;
	if (!root||!key) return;
	while (root) {
		target = __findFirstBigger(root, key);
		if (target) {
			if (target == root->keys) root = root->children[0];
			else {
				--target;
				if (root->tree->is_equal(*target, key)) {
					ret->node = root;
					ret->key = target;
					return;
				} else root = root->children[target - root->keys + 1];
			}
		} else {
			target = root->keys + root->ksz - 1;
			if (root->tree->is_equal(*target, key)) {
				ret->node = root;
				ret->key = target;
				return;
			} else if (root->csz) root = root->children[root->csz-1];
			else root = NULL;
		}
	}
}

void __insertKeyToBNode(BNode* node, void* key, BNode* rightChd, BNode** root) 
{
	void** lastPos;
	ssize_t idx, lastIdx;
	lastPos = node->keys + node->ksz - 1;
	while (lastPos>=node->keys && node->tree->is_bigger(*lastPos, key)) {
		*(lastPos+1) = *lastPos;
		--lastPos;
	}
	*(lastPos+1) = key;
	++node->ksz;
	
	if (rightChd) {
		idx = lastPos - node->keys + 2;
		lastIdx = node->csz - 1;
		while (lastIdx >= idx) {
			node->children[lastIdx+1] = node->children[lastIdx];
			++node->children[lastIdx+1]->idx;
			lastIdx--;
		}
		node->children[lastIdx+1] = rightChd;
		rightChd->idx = lastIdx+1;
		rightChd->parent = node;
		
		if (rightChd->idx) {
			rightChd->preSib = node->children[rightChd->idx-1];
			node->children[rightChd->idx-1]->nextSib = rightChd;
		} else rightChd->preSib = NULL;
		
		if (rightChd->idx < node->csz) {
			rightChd->nextSib = node->children[rightChd->idx+1];
			node->children[rightChd->idx+1]->preSib = rightChd;
		} else rightChd->nextSib = NULL;
		
		++node->csz;
	}
	
	if (NEED_SPLIT(node)) {
		if (ALLOW_ROTATE_LEFT(node)) __leftRotateKey(node);
		else if( ALLOW_ROTATE_RIGHT(node)) __rightRotateKey(node);
		else __spliteBNode(node, root);
	}
}

void InsertKeyToBTree(BNode* root, void* key, BNode** rootp) 
{
	BNode* node = root;
	void** target = NULL;
	if (!key) return;
	if (!node) {*rootp = __newRoot(key, NULL, NULL); return;}
	while (node && node->csz) {
		target = __findFirstBigger(node, key);
		if (target) node = node->children[target - node->keys];
		else node = node->children[node->csz - 1];
	}
	
	__insertKeyToBNode(node, key, NULL, rootp);
}

void __add_key(BNode* node, BNode** root);

void __mergeNode(BNode* left, BNode* right, BNode** root)
{
	ssize_t idx = 0;
	BNode** niter = NULL;
	BNode* parent = left->parent;
	void* keyFromParent = parent->keys[left->idx];
	void** iter = parent->keys + left->idx;
	while (iter != parent->keys+parent->ksz-1) {
		*(iter) = *(iter+1);
		++iter;
	}
	--parent->ksz;
	
	niter = parent->children + right->idx;
	while (niter != parent->children + parent->csz - 1) {
		*(niter) = *(niter+1);
		(*niter)->idx -= 1;
		++niter;
	}
	--parent->csz;
	
	left->keys[left->ksz] = keyFromParent;
	++left->ksz;
	
	memcpy(left->keys+left->ksz, right->keys, right->ksz*sizeof(void*));
	left->ksz += right->ksz;
	
	while (idx < right->csz) {
		left->children[left->csz+idx] = right->children[idx];
		left->children[left->csz+idx]->parent = left;
		left->children[left->csz+idx]->idx = left->csz + idx;
		++idx;
	}
	
	if (left->csz) {
		left->children[left->csz-1]->nextSib = left->children[left->csz];
		left->children[left->csz]->preSib = left->children[left->csz-1];
	}
	
	left->csz += right->csz;
	left->nextSib = right->nextSib;
	if (left->nextSib) left->nextSib->preSib = left;
	__releaseBNode(right);
	/*
	  right = left->nextSib;
	  while (right) {--right->idx; right=right->nextSib;}
	*/
	if (!parent->ksz&&!parent->parent) {
		__releaseBNode(parent);
		left->parent = NULL;
		*root = left;
		parent = NULL;
	}
	
	if (parent && *root!=parent && NEED_KEY(parent)) __add_key(parent, root);
}

void __add_key(BNode* node, BNode** root) 
{
	if (!node->parent) {
		if (!node->ksz) {
			__releaseBNode(node);
			*root = NULL;
		}
	} else if (node->preSib && HAS_SURPLUS_KEY(node->preSib)) {
		__rightRotateKey(node->preSib);
	} else if (node->nextSib && HAS_SURPLUS_KEY(node->nextSib)) {
		__leftRotateKey(node->nextSib);
	} else {
		if (node->preSib) {
			__mergeNode(node->preSib, node, root);
		} else {
			__mergeNode(node, node->nextSib, root);
		}
	}
}

void __removeKeyFromNode(BNode* node, void** target, BNode** root) 
{
	void** iter = NULL;
	BNode* leftChd = NULL;
	BNode* rightChd = NULL;
	
	if (node->csz) {
		leftChd = node->children[target-node->keys];		
		while (leftChd->csz) leftChd = leftChd->children[leftChd->csz-1];
		if (HAS_SURPLUS_KEY(leftChd)) {
			//node->tree->copy(*target, leftChd->keys[leftChd->ksz-1]);
			*target = leftChd->keys[leftChd->ksz-1];
			__removeKeyFromNode(leftChd, leftChd->keys+leftChd->ksz-1, root);
		} else {
			rightChd = node->children[target-node->keys+1];
			while (rightChd->csz) rightChd = rightChd->children[0];
			//node->tree->copy(*target, rightChd->keys[0]);
			*target = rightChd->keys[0];
			__removeKeyFromNode(rightChd, rightChd->keys, root);
		}
	} else {
		//node->tree->release_key_fn(*target);
		iter = target;
		while (iter < node->keys+node->ksz-1) {
			*(iter) = *(iter+1);
			++iter;
		}
		--node->ksz;
		if (NEED_KEY(node)) __add_key(node, root);
	}
}

void* RemoveKeyFromBTree(BNode** root, void* key)
{
	QueryRet ret;
	void* found_key;
	ret.node = NULL;
	ret.key = NULL;
	if (!root || !key) return NULL;
	QueryKeyInBTree(*root, key, &ret);
	if (ret.node && ret.key) {
		found_key = *ret.key;
		__removeKeyFromNode(ret.node, ret.key, root);
		return found_key;
	} else return NULL;
}

void __releaseBNode(BNode* node) 
{
	kfree(node->keys);	
	kfree(node->children);
	kfree(node);
}

void ReleaseBTree(BTree root) 
{
	ssize_t idx = 0;
	if (!root) return;
	while (idx < root->csz) {
		ReleaseBTree(root->children[idx]);
		if (idx < root->ksz) root->tree->release_key_fn(root->keys[idx]);
		++idx;
	}

	__releaseBNode(root);
}

void __printBTree(BTree root, ssize_t indent) 
{
	ssize_t idx = 0;
	while (idx++ < indent) printk("\t");
	idx = 0;
	while (idx < root->ksz-1) {
		root->tree->print_fn(root->keys[idx]);
		printk(",");
		++idx;
	}
	root->tree->print_fn(root->keys[idx]);
	idx = 0;
	while (idx < root->csz) {
		printk("\n");
		__printBTree(root->children[idx], indent+1);
		++idx;
	}
}

void PrintBTree(BTree root)
{
	if (!root) return;
	__printBTree(root, 0);
	printk("\n");
}


/************************ SelBalanceTree Logic *************************/

void* GetFirstKey(SelBalanceTree* tree)
{
	void *ret = NULL;
	if (!tree) goto out;
	spin_lock(&tree->lock);
	ret = GetFirstKeyFromBTree(tree->root);
	spin_unlock(&tree->lock);
out:
	return ret;
}

void* QueryKey(SelBalanceTree* tree, void* key)
{
	QueryRet ret = {NULL, NULL};
	if (!tree || !key) goto out;
	spin_lock(&tree->lock);
	QueryKeyInBTree(tree->root, key, &ret);
	spin_unlock(&tree->lock);
out:
	return ret.key;
}

void InsertKey(SelBalanceTree* tree, void* key)
{
	if (!tree || !key) return;
	spin_lock(&tree->lock);
	InsertKeyToBTree((BTree)tree->root, key, (BTree*)&tree->root);
	if (!((BTree)tree->root)->tree) ((BTree)tree->root)->tree = tree;
	spin_unlock(&tree->lock);
}

void RemoveKey(SelBalanceTree* tree, void* key)
{
	void* ret;
	if (!tree || !key) return;
	spin_lock(&tree->lock);
	ret = RemoveKeyFromBTree((BTree*)&tree->root, key);	
	spin_unlock(&tree->lock);
	if (ret) tree->release_key_fn(ret);
}

void ReleaseTree(SelBalanceTree* tree)
{
	if (!tree) return;
	spin_lock(&tree->lock);
	ReleaseBTree(tree->root);
	spin_unlock(&tree->lock);
}

void PrintTree(SelBalanceTree* tree)
{
	spin_lock(&tree->lock);
	PrintBTree(tree->root);
	spin_unlock(&tree->lock);
}
